import os
import torch
import torch.nn as nn
import cv2
import sys
sys.path.append("/home/project/Qtrainer_slim/test/yolov4") 
from models.yolov4_csp import YoloV4CspNet

def runmodel():

    imagePath = "/home/project/QZTest/ESR/images/input/yyyyy.jpg"
    outPath = "/home/project/QZTest/ESR/images/output/yyyyy.jpg"
    # weightsPath = "/home/project/Qtrainer_slim/test/yolov4/weights/yolov4-csp.weights"
    weightsPath = "/home/project/Qtrainer_slim/test/yolov4/weights/last_299.pt"

    # loadmodel
    # netparams = torch.load(weightsPath)
    # for name,value in netparams.items():
    #     print("name",name)
    #     print("value",value.shape)

    # ckpt = torch.jit.load(weightsPath)  # load checkpoint
    ckpt = torch.load(weightsPath)  # load checkpoint
    # print("ckpt", ckpt)
    # for name,value in ckpt['model'].items():
    #     print("name",name)
    #     print("value",value.shape)
    # re = ""
    # for k,v in ckpt['model'].items():

    #     print('k',k)
    #     print('v',v)
    #     re = re + k+":"+"\n"+ str(v.shape) +"\n"
    # with open("params_yolov4.txt","w") as f:
    #     f.write(re)
    yolov4Net = YoloV4CspNet()
    

    # print('checkpoint',checkpoint )
    print('===> Load checkpoint data')

    # try:
    #     ckpt['model'] = {k: v for k, v in ckpt['model'].items() if model.state_dict()[k].numel() == v.numel()}
    #     model.load_state_dict(ckpt['model'], strict=False)
    #     save_weights(model, path=saveto, cutoff=-1)
    # except KeyError as e:
    #     print(e)

    # esrmodel = IMDN_RTC_fused(in_nc=3, nf=64, num_modules=[2,2,2], out_nc=3, upscale=4)
    # esrmodel.load_state_dict(netparams,strict=True)

    # # number of parameters
    # number_parameters = sum(map(lambda x: x.numel(), esrmodel.parameters()))
    # print('Params number: {}'.format(number_parameters))

    # # 预处理 
    # img = imagetool.preImage(imagePath)

    # # GPU

    # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # img = img.to(device)
    # esrmodel = esrmodel.to(device)
    # esrmodel.eval()
    
    # result = esrmodel(img)
    # result = torch.clamp(result,0,1)

    # afimg = imagetool.afterImage(result)
    # cv2.imwrite(outPath,afimg)


    # print("result",result)
runmodel()